from chess.pieces import ChessPiece, Pawn, Rook, Knight, Bishop, Queen, King

class ChessBoard:
    def __init__(self):
        # 初始化8x8棋盘
        self.board = [[None for _ in range(8)] for _ in range(8)]
        self.setup_board()
        
        # 记录国王位置，用于快速检查将军状态
        self.king_positions = {
            'white': (7, 4),
            'black': (0, 4)
        }
        
        # 记录吃掉的棋子
        self.captured_pieces = []
        
        # 记录移动历史
        self.move_history = []
        
        # 记录是否可以王车易位
        self.castling_rights = {
            'white': {'kingside': True, 'queenside': True},
            'black': {'kingside': True, 'queenside': True}
        }
    
    def setup_board(self):
        # 设置白方棋子
        self.board[7][0] = Rook('white', 7, 0)
        self.board[7][1] = Knight('white', 7, 1)
        self.board[7][2] = Bishop('white', 7, 2)
        self.board[7][3] = Queen('white', 7, 3)
        self.board[7][4] = King('white', 7, 4)
        self.board[7][5] = Bishop('white', 7, 5)
        self.board[7][6] = Knight('white', 7, 6)
        self.board[7][7] = Rook('white', 7, 7)
        
        for col in range(8):
            self.board[6][col] = Pawn('white', 6, col)
        
        # 设置黑方棋子
        self.board[0][0] = Rook('black', 0, 0)
        self.board[0][1] = Knight('black', 0, 1)
        self.board[0][2] = Bishop('black', 0, 2)
        self.board[0][3] = Queen('black', 0, 3)
        self.board[0][4] = King('black', 0, 4)
        self.board[0][5] = Bishop('black', 0, 5)
        self.board[0][6] = Knight('black', 0, 6)
        self.board[0][7] = Rook('black', 0, 7)
        
        for col in range(8):
            self.board[1][col] = Pawn('black', 1, col)
    
    def get_piece(self, row, col):
        if 0 <= row < 8 and 0 <= col < 8:
            return self.board[row][col]
        return None
    
    def move_piece(self, from_row, from_col, to_row, to_col):
        piece = self.board[from_row][from_col]
        
        if piece is None:
            return False
        
        # 记录移动历史
        captured_piece = self.board[to_row][to_col]
        self.move_history.append({
            'piece': piece,
            'from': (from_row, from_col),
            'to': (to_row, to_col),
            'captured': captured_piece
        })
        
        # 如果目标位置有棋子，则将其加入捕获列表
        if captured_piece:
            self.captured_pieces.append(captured_piece)
        
        # 处理王车易位
        if piece.type == 'king' and abs(from_col - to_col) == 2:
            # 王车易位
            if to_col > from_col:  # 王翼易位
                rook = self.board[from_row][7]
                self.board[from_row][5] = rook
                self.board[from_row][7] = None
                rook.row, rook.col = from_row, 5
            else:  # 后翼易位
                rook = self.board[from_row][0]
                self.board[from_row][3] = rook
                self.board[from_row][0] = None
                rook.row, rook.col = from_row, 3
        
        # 更新棋子位置
        self.board[to_row][to_col] = piece
        self.board[from_row][from_col] = None
        piece.row, piece.col = to_row, to_col
        
        # 更新王的位置
        if piece.type == 'king':
            self.king_positions[piece.color] = (to_row, to_col)
        
        # 处理兵的升变
        if piece.type == 'pawn' and (to_row == 0 or to_row == 7):
            # 默认升变为皇后
            self.board[to_row][to_col] = Queen(piece.color, to_row, to_col)
        
        # 更新王车易位权
        if piece.type == 'king':
            self.castling_rights[piece.color]['kingside'] = False
            self.castling_rights[piece.color]['queenside'] = False
        elif piece.type == 'rook':
            if from_col == 0:  # 后翼车
                self.castling_rights[piece.color]['queenside'] = False
            elif from_col == 7:  # 王翼车
                self.castling_rights[piece.color]['kingside'] = False
        
        return True
    
    def is_in_check(self, color):
        king_row, king_col = self.king_positions[color]
        opponent_color = 'black' if color == 'white' else 'white'
        
        # 检查所有对方棋子，看是否可以攻击到国王
        for row in range(8):
            for col in range(8):
                piece = self.board[row][col]
                if piece and piece.color == opponent_color:
                    moves = piece.get_moves(self)
                    if (king_row, king_col) in moves:
                        return True
        
        return False
    
    def would_be_in_check(self, from_row, from_col, to_row, to_col, color):
        # 模拟走棋来检查是否会导致被将军
        piece = self.board[from_row][from_col]
        captured_piece = self.board[to_row][to_col]
        
        # 暂存国王位置
        original_king_pos = self.king_positions[color]
        
        # 暂时移动棋子
        self.board[to_row][to_col] = piece
        self.board[from_row][from_col] = None
        
        # 如果移动的是国王，更新国王位置
        if piece and piece.type == 'king' and piece.color == color:
            self.king_positions[color] = (to_row, to_col)
        
        # 检查是否在将军状态
        in_check = self.is_in_check(color)
        
        # 恢复棋盘
        self.board[from_row][from_col] = piece
        self.board[to_row][to_col] = captured_piece
        
        # 恢复国王位置
        self.king_positions[color] = original_king_pos
        
        return in_check
    
    def get_valid_moves(self, row, col):
        piece = self.board[row][col]
        if not piece:
            return []
        
        # 获取棋子的所有可能移动
        moves = piece.get_moves(self)
        
        # 过滤掉会导致自己被将军的移动
        valid_moves = []
        for move_row, move_col in moves:
            if not self.would_be_in_check(row, col, move_row, move_col, piece.color):
                valid_moves.append((move_row, move_col))
        
        # 处理王车易位
        if piece.type == 'king' and not self.is_in_check(piece.color):
            # 添加王车易位的可能移动
            if self.castling_rights[piece.color]['kingside']:
                if (self.board[row][5] is None and 
                    self.board[row][6] is None and
                    not self.would_be_in_check(row, col, row, 5, piece.color) and
                    not self.would_be_in_check(row, col, row, 6, piece.color)):
                    valid_moves.append((row, 6))
            
            if self.castling_rights[piece.color]['queenside']:
                if (self.board[row][1] is None and 
                    self.board[row][2] is None and 
                    self.board[row][3] is None and
                    not self.would_be_in_check(row, col, row, 3, piece.color) and
                    not self.would_be_in_check(row, col, row, 2, piece.color)):
                    valid_moves.append((row, 2))
        
        return valid_moves
    
    def is_checkmate(self, color):
        # 如果不在将军状态，肯定不是将死
        if not self.is_in_check(color):
            return False
        
        # 检查是否有棋子可以移动来解除将军
        for row in range(8):
            for col in range(8):
                piece = self.board[row][col]
                if piece and piece.color == color:
                    valid_moves = self.get_valid_moves(row, col)
                    if valid_moves:
                        return False
        
        # 如果没有有效移动，则是将死
        return True
    
    def is_stalemate(self, color):
        # 如果在将军状态，不是逼和
        if self.is_in_check(color):
            return False
        
        # 检查是否有棋子可以移动
        for row in range(8):
            for col in range(8):
                piece = self.board[row][col]
                if piece and piece.color == color:
                    valid_moves = self.get_valid_moves(row, col)
                    if valid_moves:
                        return False
        
        # 如果没有有效移动，且不在将军状态，则是逼和
        return True 